import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable

class Fashionnet(nn.Module):
    def __init__(self):
        super(Fashionnet, self).__init__()
        # the input shape is (512, 512, 3)
        self.conv1_1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.relu1_1 = nn.ReLU(inplace=True)
        self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
        self.relu1_2 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(2)
        
        self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.relu2_1 = nn.ReLU(inplace=True)
        self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, padding=1)
        self.relu2_2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(2)

        self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.relu3_1 = nn.ReLU(inplace=True)
        self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
        self.relu3_2 = nn.ReLU(inplace=True)
        self.conv3_3 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
        self.relu3_3 = nn.ReLU(inplace=True)
        self.pool3 = nn.MaxPool2d(2)

        self.conv4_1 = nn.Conv2d(256, 512, kernel_size=3, padding=1)
        self.relu4_1 = nn.ReLU(inplace=True)
        self.conv4_2 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.relu4_2 = nn.ReLU(inplace=True)
        self.conv4_3 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.relu4_3 = nn.ReLU(inplace=True)

    # Pose
        self.pool_pose_1 = nn.MaxPool2d(2)
        self.conv_pose_1 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.relu_pose_1 = nn.ReLU(inplace=True)
        self.conv_pose_2 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.relu_pose_2 = nn.ReLU(inplace=True)
        self.conv_pose_3 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.relu_pose_3 = nn.ReLU(inplace=True)
        self.pool_pose_2 = nn.MaxPool2d(2)
        
        # the remain feature size can be calculated by 512/(2^5)
        # as kernel_size=3, padding=1, only maxpooling will reduce the size
        self.fc6_pose = nn.Linear(512*16*16, 1024)
        self.relu_pose_4 = nn.ReLU(inplace = True)
        self.fc7_pose = nn.Linear(1024, 1024)
        self.relu_pose_5 = nn.ReLU(inplace = True)
        self.category = nn.Linear(1024, 1)
        self.location = nn.Linear(1024, 48)
        self.vis = nn.Linear(1024,24)
        self.softmax_pose = nn.Softmax()

    def num_flat_feature(self, x):
        size = x.size()[1:]
        num_feature = 1
        for s in size:
            num_features *= s
        print(num_features)
        return num_features
    
    def forward(self, x):
        x = self.relu1_1(self.conv1_1(x))
        x = self.relu1_2(self.conv1_2(x))
        x = self.pool1(x)

        x = self.relu2_1(self.conv2_1(x))
        x = self.relu2_2(self.conv2_2(x))
        x = self.pool2(x)

        x = self.relu3_1(self.conv3_1(x))
        x = self.relu3_2(self.conv3_2(x))
        x = self.relu3_3(self.conv3_3(x))
        x = self.pool3(x)

        x = self.relu4_1(self.conv4_1(x))
        x = self.relu4_2(self.conv4_2(x))
        x = self.relu4_3(self.conv4_3(x))

        # pose
        p = self.pool_pose_1(x)
        p = self.relu_pose_1(self.conv_pose_1(p))
        p = self.relu_pose_2(self.conv_pose_2(p))
        p = self.relu_pose_3(self.conv_pose_3(p))
        p = self.pool_pose_2(p)
        p = p.view(-1, 512*16*16)
        
        p = self.relu_pose_4(self.fc6_pose(p)) 
        p = self.relu_pose_5(self.fc7_pose(p))
        category = self.category(p)
        loc = self.location(p)
        vis = self.softmax_pose(self.vis(p))
        loc = loc.view(1,24,2)
        vis = vis.view(1,24,1)
        
        outputs = [category, loc, vis]
        
        return outputs


    
 #   def get_rois(loc):
 #       pre_loc = loc.view(24, 2)
 #       wh_size = torch.ones(24, 2)*3
 #       rois = torch.cat([pre_loc,wh_size], 1)
 #       rois = torch.unsqueeze(rois, 0)
 #       return rois

if __name__=='__main__':

    net = Fashionnet()
    print(net)
    
    params = list(net.parameters())
    k = 0
    for i in params:
        l = 1

        print('the structure of this layer', str(list(i.size())))

        for j in i.size():
            l *= j
        print ('the sum of params:', str(l))
        k = k + l
    
    print('finnal number of params:', str(k))
    












